{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "aaea04cfd3f16454",
   "metadata": {},
   "source": [
    "# RAG with Elastic and Llama3 using Llamaindex\n",
    "\n",
    "This interactive notebook uses `Llamaindex` to process fictional workplace documents and uses `Llama3` running locally using `Ollama` to transform these documents into embeddings and store them into `Elasticsearch`. We then ask a question, retrieve the relevant documents from `Elasticsearch` and use `Llama3` to provide a response. \n",
    "\n",
    "**_Note_** : _Llama3 is expected to be running using `Ollama` on the same machine where you will be running this notebook._"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75b9e790a7e535b0",
   "metadata": {},
   "source": [
    "## Requirements\n",
    "\n",
    "For this example, you will need:\n",
    "\n",
    "- An Elastic deployment\n",
    "  - We'll be using [Elastic Cloud](https://www.elastic.co/guide/en/cloud/current/ec-getting-started.html) for this example (available with a [free trial](https://cloud.elastic.co/registration?onboarding_token=vectorsearch&utm_source=github&utm_content=elasticsearch-labs-notebook))\n",
    "  - For LLM we will be using [Ollama](https://ollama.com/) and [Llama3](https://ollama.com/library/llama3) configured locally.  \n",
    "\n",
    "### Use Elastic Cloud\n",
    "\n",
    "If you don't have an Elastic Cloud deployment, follow these steps to create one.\n",
    "\n",
    "1. Go to [Elastic cloud Registration](https://cloud.elastic.co/registration?onboarding_token=vectorsearch&utm_source=github&utm_content=elasticsearch-labs-notebook) and sign up for a free trial\n",
    "2. Select **Create Deployment** and follow the instructions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89777ad31dbed1fe",
   "metadata": {},
   "source": [
    "## Install required dependencies for LlamaIndex and Elasticsearch\n",
    "\n",
    "First we install the packages we need for this example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c33613aab394f7d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index llama-index-cli llama-index-core llama-index-embeddings-elasticsearch llama-index-embeddings-ollama llama-index-legacy llama-index-llms-ollama llama-index-readers-elasticsearch llama-index-readers-file llama-index-vector-stores-elasticsearch llamaindex-py-client"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac6d83485e810a2",
   "metadata": {},
   "source": [
    "## Import packages\n",
    "Next we import the required packages as required. The imports are placed in the cells as required."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8fa1c4a9752298f9",
   "metadata": {},
   "source": [
    "## Prompt user to provide Cloud ID and API Key\n",
    "We now prompt the user to provide us Cloud ID and API Key using `getpass`. We get these details from the deployment. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5093cfb9b0688aa1",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-06-04T16:39:54.371586Z",
     "start_time": "2024-06-04T16:39:45.628377Z"
    }
   },
   "outputs": [],
   "source": [
    "from getpass import getpass\n",
    "\n",
    "\n",
    "# https://www.elastic.co/search-labs/tutorials/install-elasticsearch/elastic-cloud#finding-your-cloud-id\n",
    "ELASTIC_CLOUD_ID = getpass(\"Elastic Cloud ID: \")\n",
    "\n",
    "# https://www.elastic.co/search-labs/tutorials/install-elasticsearch/elastic-cloud#creating-an-api-key\n",
    "ELASTIC_API_KEY = getpass(\"Elastic Api Key: \")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24b1abd62ac05d07",
   "metadata": {},
   "source": [
    "## Prepare documents for chunking and ingestion\n",
    "We now prepare the data to be in the [Document](https://docs.llamaindex.ai/en/stable/module_guides/loading/documents_and_nodes/) type for processing using [Llamaindex](https://docs.llamaindex.ai/en/stable/) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a7d3aa4eb6faaa2d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-06-04T17:24:13.041109Z",
     "start_time": "2024-06-04T17:24:12.826203Z"
    }
   },
   "outputs": [],
   "source": [
    "import json\n",
    "from urllib.request import urlopen\n",
    "from llama_index.core import Document\n",
    "\n",
    "url = \"https://raw.githubusercontent.com/elastic/elasticsearch-labs/main/datasets/workplace-documents.json\"\n",
    "\n",
    "response = urlopen(url)\n",
    "workplace_docs = json.loads(response.read())\n",
    "\n",
    "# Building Document required by LlamaIndex.\n",
    "documents = [\n",
    "    Document(\n",
    "        text=doc[\"content\"],\n",
    "        metadata={\n",
    "            \"name\": doc[\"name\"],\n",
    "            \"summary\": doc[\"summary\"],\n",
    "            \"rolePermissions\": doc[\"rolePermissions\"],\n",
    "        },\n",
    "    )\n",
    "    for doc in workplace_docs\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d39dd6e65a8ea3df",
   "metadata": {},
   "source": [
    "## Define Elasticsearch and ingest pipeline in LlamaIndex for document processing. Use Llama3 for generating embeddings.\n",
    "We now define the `Elasticsearchstore` with the required index name, the text field and its associated embeddings. We use `Llama3` to generate the embeddings. We will be running Semantic search on the index to find documents relevant to the query posed by the user. We will use the `SentenceSplitter` provided by `Llamaindex` to chunk the documents. All this is run as part of an `IngestionPipeline` provided by the `Llamaindex` framework."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "5abfbf4f8a1b5a35",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-06-04T17:24:40.230008Z",
     "start_time": "2024-06-04T17:24:15.533532Z"
    }
   },
   "outputs": [],
   "source": [
    "from llama_index.core.node_parser import SentenceSplitter\n",
    "from llama_index.core.ingestion import IngestionPipeline\n",
    "from llama_index.embeddings.ollama import OllamaEmbedding\n",
    "from llama_index.vector_stores.elasticsearch import ElasticsearchStore\n",
    "\n",
    "es_vector_store = ElasticsearchStore(\n",
    "    index_name=\"workplace_index\",\n",
    "    vector_field=\"content_vector\",\n",
    "    text_field=\"content\",\n",
    "    es_cloud_id=ELASTIC_CLOUD_ID,\n",
    "    es_api_key=ELASTIC_API_KEY,\n",
    ")\n",
    "# Embedding Model to do local embedding using Ollama.\n",
    "ollama_embedding = OllamaEmbedding(\"llama3\")\n",
    "# LlamaIndex Pipeline configured to take care of chunking, embedding\n",
    "# and storing the embeddings in the vector store.\n",
    "pipeline = IngestionPipeline(\n",
    "    transformations=[\n",
    "        SentenceSplitter(chunk_size=512, chunk_overlap=100),\n",
    "        ollama_embedding,\n",
    "    ],\n",
    "    vector_store=es_vector_store,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b343865fcd1a742",
   "metadata": {},
   "source": [
    "## Execute pipeline \n",
    "This will chunk the data, generate embeddings using `Llama3` and ingest into `Elasticsearch` index, with embeddings in a `dense` vector field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b498e38cbc25453",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline.run(show_progress=True, documents=documents)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "301058b56643d4c5",
   "metadata": {},
   "source": [
    "The embeddings are stored in a dense vector field of dimension `4096`. The dimension size comes from the size of the embeddings generated from `Llama3`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee297886dd4e6396",
   "metadata": {},
   "source": [
    "## Define LLM settings. \n",
    "This connects to your local LLM. Please refer to https://ollama.com/library/llama3 for details on steps to run Llama3 locally. \n",
    "\n",
    "_If you have sufficient resources (atleast >64 GB Ram and GPU available) then you could try the 70B parameter version of Llama3_ "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "eaee053103a55f39",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.ollama import Ollama\n",
    "from llama_index.core import Settings\n",
    "\n",
    "Settings.embed_model = ollama_embedding\n",
    "local_llm = Ollama(model=\"llama3\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e22ed3f0b11c4bb5",
   "metadata": {},
   "source": [
    "### Setup Semantic search and integrate with Llama3. \n",
    "We now configure `Elasticsearch` as the vector store for the `Llamaindex` query engine. The query engine, using `Llama3` is then used to answer your questions with contextually relevant data from `Elasticsearch`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22e96628ac44f20d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import VectorStoreIndex, QueryBundle\n",
    "\n",
    "index = VectorStoreIndex.from_vector_store(es_vector_store)\n",
    "query_engine = index.as_query_engine(local_llm, similarity_top_k=10)\n",
    "\n",
    "# Customer Query\n",
    "query = \"What are the organizations sales goals?\"\n",
    "bundle = QueryBundle(\n",
    "    query_str=query, embedding=Settings.embed_model.get_query_embedding(query=query)\n",
    ")\n",
    "\n",
    "response = query_engine.query(bundle)\n",
    "\n",
    "print(response.response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b240b425f3b1ceae",
   "metadata": {},
   "source": [
    "_You could now try experimenting with other questions._"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
